import base64
import json
import os
import traceback
import logging
import sys  # 导入 sys 模块

from flask_socketio import SocketIO, emit
from flask import Flask, request, jsonify, send_from_directory, render_template, send_file, g, abort
import requests
import threading
import time
import uuid
import re

from flask_cors import CORS
from flask_talisman import Talisman

import client_config
from utils.loaddata import load_client_data
from sql_options import *
from pathlib import Path
import urllib3
from client import Client
from utils.options import args_parser, get_config
import torch
from data_client import ClientData, ClientStep

logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

app = Flask(__name__, static_folder="./dist/static/",
            template_folder="./dist")
cors_origins = os.getenv('CORS_ORIGIN', '10.21.231.40')
# CORS(app, resources={r"/*": {"origins": cors_origins }})
CORS(app, resources={r"/*": {"origins": "*"}})  # 允许所有域访问所有路由
# CORS(app)# 允许所有域访问所有路由

csp = {
    'default-src': ["'self'"],
    'script-src': ["'self'", "https://buttons.github.io/", "https://cdnjs.cloudflare.com",
                   "https://kit.fontawesome.com", "'nonce-randomNonceValue'"],
    'style-src': ["'self'", "https://fonts.googleapis.com", "https://cdnjs.cloudflare.com", "'nonce-randomNonceValue'"],
    'img-src': ["'self'", "data:"],
    'font-src': ["'self'", "https://fonts.gstatic.com", "https://cdnjs.cloudflare.com", "https://ka-f.fontawesome.com",
                 "data:"],
    # 'frame-ancestors': ["'self'",""],
    # 'frame-ancestors': ["'self'", ""],
    # 'connect-src': ["'self'", "https://ka-f.fontawesome.com"],
    # 'form-action': ["'self'" ],
    # 'object-src': ["'none'"],
    # 'base-uri': ["'self'"],
}

# talisman = Talisman(app, content_security_policy=csp, force_https=False)
# socketio = SocketIO(app, cors_allowed_origins="*")

client_id = None
# client_id = 'test'
temp_client_id = None
cache_data = {"client": None}
client_server = None
args = get_config()
# conn = create_connection()
# if conn:
#     create_client_data_table(conn)

# 设置数据集目录
UPLOAD_FOLDER = 'data/train/'
# 确保上传目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)


@app.errorhandler(Exception)
def handle_exception(e):
    return {"error": str(e)}, 500


@app.route("/init", methods=['GET', 'POST'])
def init():
    global client_server
    # 连接数据库
    conn = create_connection()
    # 更新多个值
    args.update({
        'device': torch.device(
            'cuda:{}'.format(args["gpu"]) if torch.cuda.is_available() and args["gpu"] != -1 else 'cpu'),
        'method': 'VeriTrac',
        'global_ep': 1
    })
    print(args["device"])
    # args = args_parser()
    # args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
    # args.method = 'VeriTrac'
    # args.global_ep = 1
    client_server = Client(args)
    headers = {'Client-ID': client_id}
    response_trace = requests.post(
        f'{client_config.SERVER_TRACE_HOST}/get_client_params',
        headers=headers, verify=False
    )
    data = response_trace.json()

    params = data.get('params')
    print("client params start")
    client_server.set_params(params, client_id, conn)
    conn.close()
    print("client params end")

    script_dir = Path(__file__).parent.resolve()
    print(script_dir)
    # 默认数据集路径
    default_dir = script_dir / 'data/train/default_data'
    print(default_dir)
    # 客户端特定数据集路径
    client_dir = script_dir / f'data/train/{client_id}'
    print(client_dir)
    # 获取默认数据集目录下的所有文件和文件夹
    default_datasets = get_sorted_filenames(directory=default_dir)
    # default_datasets = [item.name for item in default_dir.iterdir() if item.is_file() or item.is_dir()]
    # 获取客户端特定数据集目录下的所有文件和文件夹
    # 检查文件夹是否存在
    if not client_dir.exists():
        print("文件夹不存在")
        client_datasets = []
    else:
        client_datasets = get_sorted_filenames(directory=client_dir)
        # client_datasets = [item.name for item in client_dir.iterdir() if item.is_file() or item.is_dir()]
        print("文件夹中的内容:", client_datasets)
    # 合并两个数据集列表，避免重复
    # combined_datasets = list(set(default_datasets + client_datasets))  # 使用 set 去重
    combined_datasets = client_datasets + default_datasets
    # 获取两个目录下的文件信息
    default_files_info = get_files_info(directory=default_dir, isdefault=True)
    client_files_info = get_files_info(directory=client_dir)
    # 合并两个数据集的文件信息
    # combined_files_info = default_files_info + client_files_info
    # print(combined_files_info)
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("agg_verify", True)
        client_data.set_user_data("dataset", combined_datasets[0])
        client_data.set_user_data("datasets", combined_datasets)
        client_data.set_user_data("default_datasets", default_datasets)
        client_data.set_user_data("client_datasets", client_datasets)
        # client_data.set_user_data("manage_datasets", combined_files_info)
        client_data.set_user_data("manage_datasets", client_files_info)
        client_data.set_step(ClientStep.INIT_DONE.value)
    return jsonify({"datasets": combined_datasets, "client_id": client_id})

@app.route("/update_manage_datasets", methods=['GET', 'POST'])
def update_manage_datasets():
    script_dir = Path(__file__).parent.resolve()
    print(script_dir)
    # 默认数据集路径
    default_dir = script_dir / 'data/train/default_data'
    print(default_dir)
    # 客户端特定数据集路径
    client_dir = script_dir / f'data/train/{temp_client_id}'
    print(client_dir)
    # 获取默认数据集目录下的所有文件和文件夹
    default_datasets = get_sorted_filenames(directory=default_dir)
    # default_datasets = [item.name for item in default_dir.iterdir() if item.is_file() or item.is_dir()]
    # 获取客户端特定数据集目录下的所有文件和文件夹
    # 检查文件夹是否存在
    if not client_dir.exists():
        print("文件夹不存在")
        client_datasets = []
    else:
        client_datasets = get_sorted_filenames(directory=client_dir)
        # client_datasets = [item.name for item in client_dir.iterdir() if item.is_file() or item.is_dir()]
        print("文件夹中的内容:", client_datasets)
    # 合并两个数据集列表，避免重复
    # combined_datasets = list(set(default_datasets + client_datasets))  # 使用 set 去重
    combined_datasets = client_datasets + default_datasets
    # 获取两个目录下的文件信息
    default_files_info = get_files_info(directory=default_dir, isdefault=True)
    client_files_info = get_files_info(directory=client_dir)
    # 合并两个数据集的文件信息
    # combined_files_info = default_files_info + client_files_info
    # print(combined_files_info)

    client_data: ClientData = cache_data.get("client")
    if not client_data:
        client_data = ClientData("", "", "")
    # client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("dataset", combined_datasets[0])
        client_data.set_user_data("datasets", combined_datasets)
        client_data.set_user_data("default_datasets", default_datasets)
        client_data.set_user_data("client_datasets", client_datasets)
        # client_data.set_user_data("manage_datasets", combined_files_info)
        client_data.set_user_data("manage_datasets", client_files_info)
        cache_data["client"] = client_data
    return jsonify({"manage_datasets": client_files_info})
    # return jsonify({"manage_datasets": combined_files_info})


def get_sorted_filenames(directory):
    """
    获取文件夹内的文件名，按照文件创建时间排序（时间越近越靠前）
    """
    if not directory.exists() or not directory.is_dir():
        return []  # 如果目录不存在或不是文件夹，返回空列表

    # 获取文件列表，包含文件名和创建时间戳
    files_info = [
        {'filename': file.name, 'creation_time': file.stat().st_ctime}
        for file in directory.iterdir() if file.is_file()
    ]

    # 按创建时间排序，时间近的在前
    files_info.sort(key=lambda x: x['creation_time'], reverse=True)

    # 提取排序后的文件名
    sorted_filenames = [file['filename'] for file in files_info]

    return sorted_filenames


# 获取默认数据集的文件和创建时间
def get_files_info(directory, isdefault=False):
    if not directory.exists() or not directory.is_dir():
        return []  # 如果目录不存在或不是文件夹，返回空列表
    files_info = []
    for file in directory.iterdir():
        if file.is_file():
            if isdefault:
                # 获取文件的创建时间，并使用 strftime 自定义格式
                creation_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(file.stat().st_ctime))
                files_info.append({
                    'dataset_name': file.name,
                    'upload_time': creation_time,
                    'isdefault': True,
                    'creation_time': file.stat().st_ctime  # 保留时间戳方便排序
                })
            else:
                # 获取文件的创建时间，并使用 strftime 自定义格式
                creation_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(file.stat().st_ctime))
                files_info.append({
                    'dataset_name': file.name,
                    'upload_time': creation_time,
                    'isdefault': False,
                    'creation_time': file.stat().st_ctime  # 保留时间戳方便排序
                })
    # 按文件创建时间排序
    # files_info.sort(key=lambda x: x['upload_time'])
    # files_info.sort(key=lambda x: time.strptime(x['upload_time'], '%Y-%m-%d %H:%M:%S'))
    files_info.sort(key=lambda x: x['creation_time'], reverse=True)

    return files_info


# # 上传数据集
@app.route('/upload_dataset', methods=['POST'])
def upload_dataset():
    # 检查请求中是否有文件部分
    if 'files[]' not in request.files:
        return jsonify({'error': '没有上传文件'}), 400
    files = request.files.getlist('files[]')
    if len(files) == 0:
        return jsonify({'error': '没有选择文件'}), 400

    # 保存文件
    for file in files:
        if file.filename == '':
            continue  # 跳过没有文件名的文件
        # 检查文件类型是否为.npz
        if not file.filename.endswith('.npz'):
            return jsonify({'error': '文件类型不正确，只允许.npz文件'}), 400

        # 获取文件名和构造保存路径
        default_dir = os.path.join('data/train/default_data', file.filename)
        local_dataset = file.filename
        script_dir = Path(__file__).parent.resolve()
        if client_id:
            client_dir = os.path.join('data/train', client_id, file.filename)
            # client_dir = script_dir / f'data/train/{client_id}/{local_dataset}'
        else:
            client_dir = os.path.join('data/train', temp_client_id, file.filename)
        # 检查是否存在重名文件
        if os.path.exists(default_dir) or os.path.exists(client_dir):
            return jsonify({'error': '文件已存在'}), 400
        # 确保保存路径存在
        os.makedirs(os.path.dirname(client_dir), exist_ok=True)  # 创建目录（如果需要）
        # 保存文件到指定目录
        file.save(client_dir)

        # 校验数据集是否能够成功加载
        try:
            # 尝试加载数据集
            if client_id:
                dataset = load_client_data(is_train=True, datafile=local_dataset, client_id=client_id)
            else:
                dataset = load_client_data(is_train=True, datafile=local_dataset, client_id=temp_client_id)
            # 如果加载成功，更新数据
            # 重新加载数据集
            script_dir = Path(__file__).parent.resolve()
            default_dir = script_dir / 'data/train/default_data'
            if client_id:
                client_dir = script_dir / f'data/train/{client_id}'
            else:
                client_dir = script_dir / f'data/train/{temp_client_id}'
            default_datasets = get_sorted_filenames(directory=default_dir)
            # default_datasets = [item.name for item in default_dir.iterdir() if item.is_file() or item.is_dir()]
            client_datasets = get_sorted_filenames(directory=client_dir)
            # client_datasets = [item.name for item in client_dir.iterdir() if item.is_file() or item.is_dir()]
            # combined_datasets = list(set(default_datasets + client_datasets))  # 使用 set 去重
            combined_datasets = client_datasets + default_datasets
            default_files_info = get_files_info(directory=default_dir, isdefault=True)
            client_files_info = get_files_info(directory=client_dir)
            # 更新客户端数据
            client_data: ClientData = cache_data.get("client")
            if client_data:
                client_data.set_user_data("dataset", combined_datasets[0])
                client_data.set_user_data("datasets", combined_datasets)
                client_data.set_user_data("default_datasets", default_datasets)
                client_data.set_user_data("client_datasets", client_datasets)
                # client_data.set_user_data("manage_datasets", combined_files_info)
                client_data.set_user_data("manage_datasets", client_files_info)
            return jsonify({'message': '文件上传并验证成功'}), 200
        except Exception as e:
            # 如果加载数据集失败，删除上传的文件并返回错误
            logging.error(f"数据集加载失败: {str(e)}")
            os.remove(client_dir)  # 删除上传的文件
            return jsonify({'error': '数据集不适用于本任务数据集加载方式，请上传符合要求的数据集'}), 400

    return jsonify({'error': '文件上传失败'}), 400
# @app.route('/upload_dataset', methods=['POST'])
# def upload_dataset():
#     # 检查请求中是否有文件部分
#     if 'files[]' not in request.files:
#         return jsonify({'error': '没有上传文件'}), 400
#     files = request.files.getlist('files[]')
#     if len(files) == 0:
#         return jsonify({'error': '没有选择文件'}), 400
#     # 保存文件
#     for file in files:
#         if file.filename == '':
#             continue  # 跳过没有文件名的文件
#         # 检查文件类型是否为.npz
#         if not file.filename.endswith('.npz'):
#             return jsonify({'error': '文件类型不正确，只允许.npz文件'}), 400
#         # 构造文件的保存路径
#         default_dir = os.path.join('data/train/default_data', file.filename)
#         print("temp_client_id",temp_client_id)
#         if client_id:
#             client_dir = os.path.join('data/train', client_id, file.filename)
#         else:
#             client_dir = os.path.join('data/train', temp_client_id, file.filename)
#         # 检查是否存在重名文件
#         if os.path.exists(default_dir) or os.path.exists(client_dir):
#             return jsonify({'error': '文件已存在'}), 400
#
#         # 确保保存路径存在
#         # os.makedirs(os.path.dirname(default_dir), exist_ok=True)
#         os.makedirs(os.path.dirname(client_dir), exist_ok=True)  # 创建目录（如果需要）
#
#         # 保存文件到指定目录
#         # file.save(default_dir)
#         file.save(client_dir)
#
#     # 重新加载数据集
#     script_dir = Path(__file__).parent.resolve()
#     default_dir = script_dir / 'data/train/default_data'
#     if client_id:
#         client_dir = script_dir / f'data/train/{client_id}'
#     else:
#         client_dir = script_dir / f'data/train/{temp_client_id}'
#
#     default_datasets = [item.name for item in default_dir.iterdir() if item.is_file() or item.is_dir()]
#     client_datasets = [item.name for item in client_dir.iterdir() if item.is_file() or item.is_dir()]
#     # combined_datasets = list(set(default_datasets + client_datasets))  # 使用 set 去重
#     combined_datasets = client_datasets + default_datasets
#     default_files_info = get_files_info(directory=default_dir, isdefault=True)
#     client_files_info = get_files_info(directory=client_dir)
#
#     # combined_files_info = default_files_info + client_files_info
#     client_data: ClientData = cache_data.get("client")
#     if client_data:
#         client_data.set_user_data("datasets", combined_datasets)
#         # client_data.set_user_data("manage_datasets", combined_files_info)
#         client_data.set_user_data("manage_datasets", client_files_info)
#
#     return jsonify({'message': '文件上传成功'}), 200


# 删除数据集
@app.route('/delete_dataset', methods=['POST'])
def delete_dataset():
    data = request.get_json()
    file_name = data.get('dataset_name')

    if not file_name:
        return jsonify({'error': '文件名不能为空'}), 400

    # 默认路径和客户端路径
    default_data_dir = os.path.join('./data/train/default_data', file_name)
    if client_id:
        client_data_dir = os.path.join('./data/train', client_id, file_name)
    else:
        client_data_dir = os.path.join('./data/train', temp_client_id, file_name)
    # 遍历两个目录，检查文件是否存在
    try:
        file_deleted = False
        # 删除默认数据集目录中的文件
        if os.path.exists(default_data_dir):
            os.remove(default_data_dir)
            file_deleted = True
        # 删除客户端数据集目录中的文件
        if os.path.exists(client_data_dir):
            os.remove(client_data_dir)
            file_deleted = True
        # 如果文件在两个文件夹中都没有找到
        if not file_deleted:
            return jsonify({'error': '文件不存在'}), 404
        # 成功删除
        # 重新加载数据集
        script_dir = Path(__file__).parent.resolve()
        default_dir = script_dir / 'data/train/default_data'
        client_dir = script_dir / f'data/train/{temp_client_id}'
        default_datasets = get_sorted_filenames(directory=default_dir)
        # default_datasets = [item.name for item in default_dir.iterdir() if item.is_file() or item.is_dir()]
        client_datasets = get_sorted_filenames(directory=client_dir)
        # client_datasets = [item.name for item in client_dir.iterdir() if item.is_file() or item.is_dir()]
        # combined_datasets = list(set(default_datasets + client_datasets))  # 使用 set 去重
        combined_datasets = client_datasets + default_datasets
        default_files_info = get_files_info(directory=default_dir, isdefault=True)
        client_files_info = get_files_info(directory=client_dir)

        # combined_files_info = default_files_info + client_files_info
        client_data: ClientData = cache_data.get("client")
        if client_data:
            client_data.set_user_data("dataset", combined_datasets[0])
            client_data.set_user_data("datasets", combined_datasets)
            client_data.set_user_data("default_datasets", default_datasets)
            client_data.set_user_data("client_datasets", client_datasets)
            # client_data.set_user_data("manage_datasets", combined_files_info)
            client_data.set_user_data("manage_datasets", client_files_info)
        return jsonify({'success': True}), 200

    except Exception as e:
        return jsonify({'error': str(e)}), 500


# @app.route('/upload_files', methods=['POST'])
# def upload_files():
#     # 检查请求中是否有文件部分
#     if 'files[]' not in request.files:
#         return jsonify({'error': '没有上传文件'}), 400
#     files = request.files.getlist('files[]')
#     if len(files) == 0:
#         return jsonify({'error': '没有选择文件'}), 400
#
#     # 保存文件
#     for file in files:
#         if file.filename == '':
#             continue  # 跳过没有文件名的文件
#         script_dir = Path(__file__).parent.resolve()
#         relative_path = 'data/train/'
#         target_dir = script_dir / relative_path
#         # 确保保存路径存在
#         target_dir = os.path.join(UPLOAD_FOLDER, file.filename)
#         os.makedirs(os.path.dirname(target_dir), exist_ok=True)  # 创建目录（如果需要）
#         # 保存文件到指定目录
#         file.save(target_dir)
#
#     return jsonify({'message': '文件上传成功'}), 200


# @app.route('/setparams', methods=['POST'])
# def setparams():
#     #连接数据库
#     conn = create_connection()
#     data = request.json
#     client_server.set_params(data,client_id, conn)
#     conn.close()
#     return jsonify({'status': 'success'}), 200

@app.route('/local_train', methods=['POST'])
def local_train():
    # 连接数据库
    data = request.json
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.reset_status()
        client_data.set_user_data("setting", get_text_with_setting(data.get('local_setting')))
        client_data.set_user_data("dataset", data.get('local_dataset'))
        client_data.set_step(ClientStep.TRAIN_START.value)
    error = None
    try:
        conn = create_connection()
        client_server.local_train(local_set=data.get('local_setting'), local_dataset=data.get('local_dataset'),
                                  client_id=client_id, connection=conn)
    except Exception as e:
        print("local_train error: {}".format(e))
        error = str(e)
    finally:
        conn.close()
    # socketio.emit('local_train')
    if error: return jsonify({'status': f'local_train error:{error}'}), 500
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("status", "训练完成")
        client_data.set_user_data("agg_verify", True)
        client_data.set_user_data("setting", get_text_with_setting(data.get('local_setting')))
        client_data.set_user_data("dataset", data.get('local_dataset'))
        client_data.set_step(ClientStep.TRAIN_DONE.value)
    return jsonify({'status': 'local_train success'}), 200


def get_text_with_setting(code):
    """
    switch (setting) {
    case '默认设置':
      return 0;
    case '伪造hash':
      return 1;
    case '伪造签名1':
      return 2;
    case '伪造签名2':
      return 3;
    default:
      return -1; // 或其他合适的默认值
  }

    """
    data = {
        0: "默认设置",
        1: "伪造hash",
        2: "伪造签名1",
        3: "伪造签名2"
    }
    return data.get(int(code), -1)


@app.route('/upload', methods=['POST'])
def upload():
    # 连接数据库
    conn = create_connection()
    upload_data = get_client_data(conn, client_id)
    conn.close()
    if 'local_BFV' in upload_data and upload_data['local_BFV'] is not None:
        upload_data['local_BFV'] = base64.b64encode(upload_data['local_BFV']).decode('utf-8')
    response_trace = requests.post(
        f'{client_config.SERVER_AGG_HOST}/upload',
        json={'upload_data': upload_data}, verify=False)
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("status", "已上传")
        client_data.set_step(ClientStep.UPLOAD_DONE.value)
    return jsonify({'status': 'upload success'}), 200


@app.route('/re_train', methods=['POST'])
def re_train():
    # socketio.emit('agg_verify')
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("agg_verify", False)
        client_data.set_step(ClientStep.INIT_DONE.value)
    return jsonify({'status': 'success'}), 200


@app.route('/update', methods=['POST'])
def update():
    global update_data
    update_data = request.json
    # socketio.emit('update', {'status': True})
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("status", "已接收")
        client_data.set_step(ClientStep.UPDATE_DONE.value)
    return jsonify({'status': 'success'}), 200


@app.route('/verify', methods=['POST'])
def verify():
    # 连接数据库
    conn = create_connection()
    verify = client_server.client_verify(data=update_data, client_id=client_id, connection=conn)
    conn.close()
    print("verify result:", verify)
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("result1", "验证成功" if verify == 1 else "验证失败")
        client_data.set_step(ClientStep.VERIFY_TRUE.value if verify == 1 else ClientStep.VERIFY_FALSE.value)
    if verify == 1:
        return jsonify({'status': 'verify success', 'verify_result': True}), 200
    else:
        return jsonify({'status': 'verify failed', 'verify_result': False})


@app.route('/request_trace', methods=['POST'])
def request_trace():
    headers = {'Client-ID': client_id}
    response_trace = requests.post(
        f'{client_config.SERVER_TRACE_HOST}/trace',
        headers=headers, verify=False)
    client_data: ClientData = cache_data.get("client")
    if client_data:
        client_data.set_user_data("result2", "处理中...")
    print(f"Sent data to trace server: {response_trace.json()}")
    return jsonify({'status': 'start trace'}), 200


@app.route('/trace_result', methods=['POST'])
def trace_result():
    data = request.get_json()
    print("trace result", data.get('unhonest_participant'))
    # socketio.emit('return_result', {'unhonest_participant': data.get('unhonest_participant')})
    client_data: ClientData = cache_data.get("client")
    # 提取用户名并用逗号分隔
    usernames = ','.join(data.get('unhonest_participant'))
    #
    print(usernames)  # 输出
    if client_data:
        client_data.set_user_data("result2", usernames)
        client_data.set_step(ClientStep.TRACE_DONE.value)
    return jsonify({'status': 'success'}), 200


@app.route("/Getmodel", methods=['POST', 'GET'])
def getmodel():
    model_path = f'./model/model.pth'
    print("MODEL_PATH:", model_path)
    try:
        return send_file(model_path, as_attachment=True)
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/Getfiles')
def get_files():
    file = request.args.get('file')
    if file == "guide":
        file_path = f"demo/model.pth的使用说明.docx"
        try:
            return send_file(file_path, as_attachment=True)
        except Exception as e:
            traceback.print_exc()
            return jsonify({'error': str(e)}), 500
    elif file == "demo":
        file_path = f"demo/demo.py"
        try:
            return send_file(file_path, as_attachment=True)
        except Exception as e:
            traceback.print_exc()
            return jsonify({'error': str(e)}), 500


# 注销客户端
@app.route('/logout', methods=['POST'])
def logout():
    p_client_id = request.args.get("client_id")
    server_trace_url = f'{client_config.SERVER_TRACE_HOST}/delete_client'
    server_agg_url = f'{client_config.SERVER_AGG_HOST}/delete_client'
    global cache_data, client_id, client_server
    d_client_id = p_client_id if p_client_id else client_id
    try:
        trace_response = requests.delete(server_trace_url, json={'client_id': d_client_id}, verify=False)
        agg_response = requests.delete(server_agg_url, json={'client_id': d_client_id})

        # 删除客户端数据表
        delete_client_data(d_client_id)

        if trace_response.status_code == 200 and agg_response.status_code == 200:
            return jsonify({'message': 'Client logged out successfully'}), 200
        else:
            error_message = []
            if trace_response.status_code != 200:
                error_message.append(f"Server Trace Error: {trace_response.text}")
            if agg_response.status_code != 200:
                error_message.append(f"Server Aggregate Error: {agg_response.text}")
            return jsonify(
                {'error': 'Failed to delete client data from one or both servers', 'details': error_message}), 500

    except requests.exceptions.RequestException as e:
        print("Request error:", str(e))
        return jsonify({'error': str(e)}), 500
    finally:
        cache_data["client"] = None
        client_id = None
        client_server = None


@app.route('/data', methods=['POST'])
def send_data(SSO_session_data):
    data = {
        'message': 'Hello from client',
        'client_ip': client_config.SERVER_CLIENT_HOST,
        # 'client_port':  client_config.SERVER_Client_PORT,
        'userid': int(SSO_session_data["userId"]),
        'username': str(SSO_session_data["username"]),
    }
    headers = {'Client-ID': client_id}
    print("headers: ", headers)
    try:
        response_trace = requests.post(
            f'{client_config.SERVER_TRACE_HOST}/data',
            json=data, headers=headers, verify=False)
        print(f"Sent data to trace server: {response_trace.json()}")

        response_agg = requests.post(f'{client_config.SERVER_AGG_HOST}/data',
                                     json=data, headers=headers, verify=False)
        print(f"Sent data to agg server: {response_agg.json()}")

    except requests.exceptions.RequestException as e:
        print(f"Failed to send data: {e}")


@app.route('/ping', methods=['GET'])
def ping():
    return jsonify({'status': 'alive', 'client_id': client_id}), 200


def run_flask():
    app.config['DEBUG'] = False
    app.run(host='0.0.0.0', port=client_config.SERVER_PORT)
    # socketio.run(app, host='0.0.0.0', port=5003, allow_unsafe_werkzeug=True)


# WebSocket 连接事件
# @socketio.on('connect')
# def handle_connect():
#     print('Client connected')


# @socketio.on('disconnect')
# def handle_disconnect():
#     print('Client disconnected')


# @app.route('/connect_to_trace_and_agg', methods=['POST'])
def connect_to_trace_and_agg(user_info):
    global cache_data
    # 连接数据库
    conn = create_connection()
    client_id = user_info["client_id"]
    userid = user_info["user_id"]
    username = user_info["username"]
    result = get_client_id(conn, client_id)
    conn.close()
    if result:
        print("Client ID:", client_id)
        # client_id = result[0]
    else:
        print("No client_id found and client_id is None")
        # client_id = str(uuid.uuid4())
        # print(client_id)
        conn = create_connection()
        insert_or_update_client_data(conn, client_id=client_id, userid=userid, username=username)
        conn.close()

    headers = {'Client-ID': client_id}
    try_count = 0
    while True:
        try:
            response_trace = requests.get(
                f'{client_config.SERVER_TRACE_HOST}/ping',
                headers=headers, verify=False)
            print(f"Connected to trace server: {response_trace.json()}")

            response_agg = requests.get(f'{client_config.SERVER_AGG_HOST}/ping',
                                        headers=headers, verify=False)
            print(f"Connected to agg server: {response_agg.json()}")
            send_data({
                "userId": userid,
                "username": username
            })
            break
        except requests.exceptions.RequestException as e:
            print("Connection failed, retrying... {}".format(e))
            time.sleep(5)
        finally:
            try_count += 1
            if try_count > 5:
                break
    # return jsonify({'status': 'success'}), 200


# 假设我们用一个全局变量来存储锁状态
# True 表示锁定，False 表示解锁
lock_status = False  # 初始状态为解锁


@app.route('/api/check_lock', methods=['GET'])
def check_lock():
    """
    检查锁的状态
    """
    global lock_status
    return jsonify({'locked': lock_status})


# 路由处理 favicon.ico 请求
@app.route('/favicon.ico')
def favicon():
    return send_from_directory(os.path.join(app.root_path, 'static'), 'favicon.png', mimetype='image/png')


@app.route('/Client')
@app.route('/')
@app.route('/<path:path>')
def serve(path='index.html'):
    user_info = paser_h_user_info(request)
    if not user_info:
        raise Exception("没有加载到用户信息")
    global temp_client_id
    temp_client_id = user_info["client_id"]
    if path != '' and os.path.exists(os.path.join(app.template_folder, path)):
        # 检查路径是否指向模板文件（如 .html 文件）
        if path.endswith('.html'):
            new_html = add_user_info(user_info, render_template(path))
            # print(new_html)
            return new_html
        else:
            # 如果路径指向静态文件，则返回静态文件
            return send_from_directory(app.template_folder, path)
    else:
        return add_user_info(user_info, render_template('index.html'))


@app.route('/login', methods=['POST'])
def login():
    trace_step = get_trace_step_status()
    print("trace_step: ", trace_step)
    if (trace_step == "step_none"):
        user_info = paser_h_user_info(request)
        if not user_info:
            raise Exception("没有加载到用户信息")
        global client_id
        client_id = user_info["client_id"]
        userid = user_info["user_id"]
        username = user_info["username"]
        firstLogin_login_id = userid
        if not cache_data["client"]:
            client_data = ClientData(client_id, userid, username)
            client_data.set_step(ClientStep.LOGIN_DONE.value)
            cache_data["client"] = client_data
            connect_to_trace_and_agg(user_info)
        else:
            client_data = cache_data["client"]
            if client_data.id != client_id:
                client_data = ClientData(client_id, userid, username)
                client_data.set_step(ClientStep.LOGIN_DONE.value)
                cache_data["client"] = client_data
                connect_to_trace_and_agg(user_info)
            firstLogin_login_id = client_data.user_id
        return jsonify({'status': 'success', "firstLogin": firstLogin_login_id}), 200
    else:
        return jsonify({'status': 'wait'}), 200


def paser_h_user_info(request):
    d_token = request.args.get("d_token")
    token = request.headers.get("X-UserInfo")
    if not token:
        token = d_token
    user_info = paser_token(token)
    if not user_info and request.method == "POST":
        json_data = request.get_json()
        if json_data and json_data.get('user'):
            user_info = json_data.get("user")
            user_info["user_id"] = user_info.get("userid")
    print("paser_h_user_info : ", user_info)
    logging.info("user_info: {}".format(user_info))
    return user_info


def paser_token(token):
    if token:
        token_data = json.loads(base64.b64decode(token).decode('utf-8'))
        client_id = str(token_data["realm"])
        userid = int(token_data["userId"])
        username = str(token_data["username"])
        return {
            "client_id": client_id,
            "user_id": userid,
            "username": username
        }
    return None


def add_user_info(user_info, html_content):
    if not user_info:
        return html_content
    element = '<user client_id="{}" userid="{}" username="{}">'.format(user_info['client_id'],
                                                                       user_info['user_id'],
                                                                       user_info['username'])
    new_content = re.sub(
        r'<user[^>]*>',  # 匹配 <user> 标签及其属性
        element,  # 替换内容
        html_content
    )
    return new_content


def get_trace_step_status():
    response = requests.get(f'{client_config.SERVER_TRACE_HOST}/trace_status?step_status=true', verify=False)
    if response.status_code == 200:
        step_status = response.json()["step_status"]
        print("Get trace step status successfully: {}".format(step_status))
        return step_status
    else:
        print("Get trace step status error: {}".format(response.text))
        return None


def get_agg_step_status():
    try:
        response = requests.get(f'{client_config.SERVER_AGG_HOST}/agg_status?step_status=true', verify=False)
        if response.status_code == 200:
            print(response.text)
            print(response.json())
            step_status = response.json()["step_status"]
            print("Get agg step status successfully: {}".format(step_status))
            return step_status
        else:
            print("Get agg step status error: {}".format(response.text))
            return None
    except Exception as e:
        print("Get agg step status: {}".format(e))


@app.route('/client_status', methods=['GET'])
def client_status():
    """
    返回客户端状态及会话缓存
    """
    trace_step = get_trace_step_status()
    agg_step = get_agg_step_status()
    user_info = paser_h_user_info(request)
    print(f"user_info: {user_info}")
    print("trace_step: ", trace_step)
    print("agg_step: ", agg_step)

    client_info: ClientData = cache_data.get("client")
    if not client_info:
        client_info = ClientData("", "", "")
    client_info.set_trace_step(trace_step)
    client_info.set_agg_step(agg_step)
    print(client_info.data())
    return client_info.data()


# 定义禁止访问的主机列表
BLOCKED_HOSTS = ['169.254.169.254', 'aws.zaproxy.org', '100.100.100.200', 'alibaba.zaproxy.org']


@app.before_request
def block_metadata_access():
    if request.host in BLOCKED_HOSTS:
        abort(403)  # 禁止访问


if __name__ == "__main__":
    # 检查和初始化数据库
    conn = create_connection()
    if conn:
        create_client_data_table(conn)
    conn.close()
    # flask_thread = threading.Thread(target=run_flask)
    # flask_thread.start()

    # connect_thread = threading.Thread(target=connect_to_trace_and_agg)
    # connect_thread.start()

    # Wait for threads to complete
    # connect_thread.join()
    # flask_thread.join()
    run_flask()

    # conn = create_connection()
    # if conn:
    #     create_client_data_table(conn)
    # app.config['DEBUG'] = False
    # # app.run(host='0.0.0.0', port=5003)
    # socketio.run(app, host='0.0.0.0', port=5003)
    # connect_thread = threading.Thread(target=connect_to_trace_and_agg)
    # connect_thread.start()
    #
    # # Wait for threads to complete
    # connect_thread.join()
    # conn.close()
